from typing import Any, Optional
from fastapi import status
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.responses import PlainTextResponse
from fastapi.encoders import jsonable_encoder

_ExcludeData = set[int | str] | dict[int | str, Any]


def create_response(
    data: Any | None = None,
    code: int = 0,
    status_code: int = 200,
    msg: str = 'Success',
    exclude: _ExcludeData | None = None,
    as_text: Optional[bool] = False,
    headers=None,
    **kwargs,
):
    if as_text:
        return PlainTextResponse(content=str(data), status_code=status_code, headers=headers, **kwargs)
    resp = {'code': code, 'msg': msg, 'data': data}
    return JSONResponse(
        status_code=status_code,
        content=jsonable_encoder(resp, exclude=exclude)
    )


def success(
    *,
    code: int = 0,
    status_code: int = 200,
    msg: str = 'Success',
    data: Any | None = None,
    exclude: _ExcludeData | None = None,
    **kwargs,
):
    """
    请求成功返回通用方法
    :param code: api状态码
    :param status_code: 返回状态码
    :param msg: 返回信息
    :param data: 返回数据
    :param exclude: 排除返回数据(data)字段
    :return:
    """
    return create_response(code=code, status_code=status_code, msg=msg, data=data, exclude=exclude, **kwargs)


def fail(
    *,
    code: int = -1,
    status_code: int = 400,
    msg: str = 'Bad Request',
    data: Any = None,
    exclude: _ExcludeData | None = None,
    **kwargs,
):
    """
    请求成功返回通用方法
    :param code: api状态码
    :param status_code: 返回状态码
    :param msg: 返回信息
    :param data: 返回数据
    :param exclude: 排除返回数据(data)字段
    :return:
    """
    return create_response(code=code, status_code=status_code, msg=msg, data=data, exclude=exclude, **kwargs)


def streaming(*, data: Any = None):
    """
    流式响应方法
    :param data: 流式传输的内容
    :return: 流式响应结果
    """
    return StreamingResponse(
        status_code=status.HTTP_200_OK,
        content=data
    )

